﻿/*
Copyright (c) 2012 Eric Lasota

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
*/

using System;
using System.Collections.Generic;

using System.Text;

namespace MPEGnet
{
    internal class RingBuffer
    {
	    byte[] _bytes;
	    uint _start;
	    uint _used;
        uint _bufferSize;

	    public RingBuffer(uint bufferSize)
	    {
		    _start = _used = 0;
            _bytes = new byte[bufferSize];
            _bufferSize = bufferSize;
	    }

        public void Write(byte[] srcBytes, int srcOffset, uint size)
	    {
		    uint startLocation = _start + _used;
		    if(startLocation >= _bufferSize)
                Array.Copy(srcBytes, srcOffset, _bytes, startLocation - _bufferSize, size);
		    else
		    {
			    uint endWritable = _bufferSize - startLocation;
			    if(endWritable >= size)
                    Array.Copy(srcBytes, srcOffset, _bytes, startLocation, size);
			    else
			    {
                    Array.Copy(srcBytes, srcOffset, _bytes, startLocation, endWritable);
                    Array.Copy(srcBytes, srcOffset + endWritable, _bytes, 0, size - endWritable);
			    }
		    }

		    _used += size;
	    }

        public uint PeekSmall(uint nBytes)
        {
            if (nBytes > _used)
                return 0;

            if(nBytes == 1)
                return _bytes[_start];

            uint endAvailable = _bufferSize - _start;

            if (endAvailable >= 4)
            {
                uint result = System.BitConverter.ToUInt32(_bytes, (int)_start);
                if(System.BitConverter.IsLittleEndian)

                    result = (
                        ((result << 24) & 0xff000000) |
                        ((result << 8) & 0xff0000) |
                        ((result >> 8) & 0xff00) |
                        ((result >> 24) & 0xff)
                        );

                result >>= ((4 - (int)nBytes) * 8);
                return result;
            }
            else if (endAvailable >= nBytes)
            {
                uint result = 0;
                for (uint i = 0; i < nBytes; i++)
                    result = ((result << 8) | _bytes[_start + i]);
                return result;
            }
            else
            {
                uint result = 0;
                for (uint i = 0; i < endAvailable; i++)
                    result = ((result << 8) | _bytes[_start + i]);
                for (uint i = 0; i < nBytes - endAvailable; i++)
                    result = ((result << 8) | _bytes[i]);

                return result;
            }
        }

        public uint Peek(byte[] buf, int bufOffset, uint size)
	    {
		    if(size > _used)
		    {
                Array.Clear(buf, bufOffset + (int)_used, (int)(size - _used));
			    size = _used;
		    }

		    uint endAvailable = _bufferSize - _start;

		    if(endAvailable >= size)
                Array.Copy(_bytes, _start, buf, bufOffset, size);
		    else
		    {
                Array.Copy(_bytes, _start, buf, bufOffset, endAvailable);
                Array.Copy(_bytes, 0, buf, bufOffset + endAvailable, size - endAvailable);
		    }

		    return size;
	    }

        public void Skip(uint size)
	    {
		    if(size > _used)
			    size = _used;
		    _start += size;
		    _used -= size;

		    if(_start >= _bufferSize)
			    _start -= _bufferSize;
	    }

        public uint Read(byte[] buf, int bufOffset, uint size)
	    {
		    uint amountRead = Peek(buf, bufOffset, size);
		    Skip(size);

		    return amountRead;
	    }

	    public uint BytesStored
	    {
            get
            {
    		    return _used;
            }
	    }

        public uint BytesFree
	    {
            get
            {
		        return _bufferSize - _used;
            }
	    }
    }

    public class VideoDecodeState : IMpegVideoIO
    {
        public PerfDebugger PDB = new PerfDebugger();

        const int DEMUX_BUFFER_SIZE = 10000;
        const int VIDEO_BUFFER_SIZE = 20000;

	    enum DemuxTarget
	    {
		    Audio,
		    Video,
		    Trash,
		    Parsing,
	    }

	    enum PS
	    {
		    PackHeader,
		    SystemHeader,
		    SystemHeaderStreams,
		    DemuxPacket,
		    Exhausted,
	    }

        RingBuffer _videoRing = new RingBuffer(VIDEO_BUFFER_SIZE);
        RingBuffer _audioRing = new RingBuffer(MP2DecodeState.FRAME_BUFFER_SIZE_BYTES + MP2DecodeState.FRAME_HEADER_SIZE_BYTES);
        RingBuffer _demuxRing = new RingBuffer(DEMUX_BUFFER_SIZE);

        byte[] _audioData = new byte[MP2DecodeState.FRAME_BUFFER_SIZE_BYTES];
	    byte[] _demuxMoveBuffer = new byte[DEMUX_BUFFER_SIZE];

        short[][] _audioSamples = new short[][]
            {
                new short[MP2DecodeState.FRAME_NUM_SAMPLES],
                new short[MP2DecodeState.FRAME_NUM_SAMPLES]
            };

	    MP2Header _audioHeader;
	    MP2DecodeState _audioDec = new MP2DecodeState();
	    VidStream _videoDec;

	    bool _parsedSystemHeader;
	    bool _flushing;
	    bool _audioAvailable;
        bool _audioExhausted;
        bool _videoInitialized;
	    bool _videoExhausted;	// If this is true, then no more video will ever be decoded
        bool _videoAvailable;
        ImageInfo _videoImageInfo;
	    bool _mustParsePacketHeaders;

	    PS _parseState;

        DemuxTarget _demuxTarget;
	    bool _noDemuxLimit;
	    uint _demuxRemaining;
        ThreadedDispatch _dispatch;

        public VideoDecodeState(int numBackgroundThreads)
        {
            _dispatch = new ThreadedDispatch(numBackgroundThreads);
		    _parseState = PS.PackHeader;
		    _demuxTarget = DemuxTarget.Parsing;
        }

        public uint BytesFree
        {
            get
            {
                return _demuxRing.BytesFree;
        	}
        }

	    public void FeedData(byte[] data, int dataOffs, uint size)
	    {
            _demuxRing.Write(data, dataOffs, size);
	    }

	    void TryParsePackHeader()
	    {
		    if(_parseState != PS.PackHeader)
			    return;

		    if(_demuxRing.BytesStored < 4)
			    return;

            uint header = _demuxRing.PeekSmall(4);

		    if(header != 0x000001ba)
		    {
			    _parseState = PS.DemuxPacket;
			    _noDemuxLimit = true;

                if (header == 0x000001b3)                        // Detect MPEG-1 video start
				    _demuxTarget = DemuxTarget.Video;
                else if ((header & 0xfff00000) == 0xfff00000)    // Detect layer 2 sync word
				    _demuxTarget = DemuxTarget.Audio;
			    else
				    _demuxTarget = DemuxTarget.Trash;


			    return;
		    }

		    // Try to parse the header info
		    if(_demuxRing.BytesStored < 24)
			    return;

            _demuxRing.Skip(12);
		
		    // We can just ignore all of this

		    // Pack:
		    //(b[4] & 0xf0) == 0x20
		    //SCR = ((b[4] & 0x7) << 30) | ((b[5]) << 22) | ((b[6] & 0xfe) << 15) | ((b[7]) << 7) | ((b[8] & 0xfe) >> 1);
		    //MuxRate = ((b[9] & 0x7f) << 15); | (b[10] << 7) | ((b[11] & 0xfe) >> 1)

		    _parseState = PS.DemuxPacket;
	    }

        void TryParseSystemHeader()
	    {
		    if(_parseState != PS.SystemHeader)
			    return;

		    // This should only ever be reachable by TryDemux setting it

            _demuxRing.Skip(12);
		    // System:
		    // RateBound = ((b[8] & 0xfe) >> 1) | (b[7] << 7) | ((b[6] & 0x7f) << 15)
		    // AudioBound = ((b[9] & 0xfc) >> 2)

		    _parseState = PS.SystemHeaderStreams;
	    }


        void TryParseStreams()
	    {
		    if(_parseState != PS.SystemHeaderStreams)
			    return;

		    if(_demuxRing.BytesStored < 3)
			    return;
		
		    uint marker = _demuxRing.PeekSmall(3);

		    if(marker == 0x000001)
		    {
			    _parseState = PS.DemuxPacket;
			    _demuxTarget = DemuxTarget.Parsing;
		    }
		    else
                _demuxRing.Skip(3);	// Skip past it
	    }

        void TryDemux()
	    {
		    if(_parseState != PS.DemuxPacket)
			    return;

		    if(_demuxTarget == DemuxTarget.Parsing)
		    {
			    if(_demuxRing.BytesStored >= 6)
			    {
                    uint headerWord = _demuxRing.PeekSmall(4);

                    if((headerWord & 0xffffff00) != 0x00000100)
                    {
                        _demuxRing.Skip(1);
                        return;
                    }

				    if(headerWord == 0x1c0)
					    _demuxTarget = IgnoreAudio ? DemuxTarget.Trash : DemuxTarget.Audio;
				    else if(headerWord == 0x1e0)
					    _demuxTarget = IgnoreVideo ? DemuxTarget.Trash : DemuxTarget.Video;
				    else if(headerWord == 0x1ba)
				    {
					    // Hit an inline pack header
					    _parseState = PS.PackHeader;
					    return;
				    }
				    else if(headerWord == 0x1bb)
				    {
					    // Hit an inline system header
					    _parseState = PS.SystemHeader;
					    return;
				    }
				    else
					    _demuxTarget = DemuxTarget.Trash;   // Padding or something

				    _mustParsePacketHeaders = true;
                    _demuxRing.Skip(4);
                    _demuxRemaining = _demuxRing.PeekSmall(2);
                    _demuxRing.Skip(2);
			    }
			    else if(_flushing)
				    _parseState = PS.Exhausted;
		    }

		    if(_demuxTarget != DemuxTarget.Parsing)
		    {
			    if(_mustParsePacketHeaders)
			    {
                    // Need to skip the packet headers
				    if(_demuxRing.BytesStored < 1)
				    {
					    if(_flushing)
						    _parseState = PS.Exhausted;
					    return;
				    }
                    uint flagLeadByte = _demuxRing.PeekSmall(1);

                    // Skip padding bytes
                    if (flagLeadByte == 0xff)
                    {
                        uint padSkip;
                        if (_demuxRing.BytesStored >= 4)
                        {
                            uint padBytes = _demuxRing.PeekSmall(4);
                            if (padBytes == 0xffffffff)
                                padSkip = 4;
                            else if ((padBytes & 0xffffff00) == 0xffffff)
                                padSkip = 3;
                            else if ((padBytes & 0xffff0000) == 0xffff)
                                padSkip = 2;
                            else
                                padSkip = 1;
                        }
                        else
                            padSkip = 1;
                        _demuxRing.Skip(padSkip);

                        if (_demuxRemaining <= padSkip)
                        {
                            // Ran out of demux and still parsing headers, just drop the packet
                            _demuxRemaining = 0;
                            _mustParsePacketHeaders = false;
                        }
                        else
                            _demuxRemaining -= padSkip;
                        return;   // Padding
                    }

				    uint skipSize = 0;
                    uint stdSize = 0;
                    uint ptsdtsSize = 0;
				    if(flagLeadByte != 0x0f)	// Skip STD/PTS/DTS
				    {
                        if (_demuxRing.BytesStored < 2)
					    {
						    if(_flushing)
							    _parseState = PS.Exhausted;
						    return;
					    }

                        uint flagBytes = _demuxRing.PeekSmall(2);
					    if((flagBytes & 0xc000) == 0x4000)
						    stdSize = 2;

                        if (_demuxRing.BytesStored < 1 + stdSize)
					    {
						    if(_flushing)
							    _parseState = PS.Exhausted;
						    return;
					    }

					    flagBytes = _demuxRing.PeekSmall(1 + stdSize);
					    if((flagBytes & 0xf0) == 0x30)
						    ptsdtsSize = 10;	// Both PTS and DTS are present
					    else if((flagBytes & 0xf0) == 0x20)
						    ptsdtsSize = 5;	// Only PTS is present

					    skipSize = stdSize + ptsdtsSize;
				    }
				    else
					    skipSize = 1;

                    if (_demuxRing.BytesStored < skipSize)
				    {
					    if(_flushing)
						    _parseState = PS.Exhausted;
					    return;
				    }
                    _demuxRing.Skip(skipSize);
				    if(_demuxRemaining < skipSize)
					    _demuxRemaining = 0;	// Ugh
				    else
					    _demuxRemaining -= skipSize;

				    _mustParsePacketHeaders = false;
			    }

			    if(!_mustParsePacketHeaders && (_demuxRemaining != 0 || _noDemuxLimit))
			    {
				    // See how much data can be moved
                    uint amount = _demuxRing.BytesStored;
				    if(amount > _demuxRemaining && !_noDemuxLimit)
					    amount = _demuxRemaining;

				    if(amount == 0)
				    {
					    // Nothing left in the demux buffer
					    if(_flushing)
						    _parseState = PS.Exhausted;
					    return;
				    }

				    if(_demuxTarget == DemuxTarget.Audio)
				    {
					    if(_audioRing.BytesFree < amount)
                            amount = _audioRing.BytesFree;

					    if(amount != 0)
					    {
						    _demuxRing.Read(_demuxMoveBuffer, 0, amount);
						    _audioRing.Write(_demuxMoveBuffer, 0, amount);
					    }
				    }
				    else if(_demuxTarget == DemuxTarget.Video)
				    {
					    if(_videoRing.BytesFree < amount)
                            amount = _videoRing.BytesFree;

					    if(amount != 0)
					    {
                            _demuxRing.Read(_demuxMoveBuffer, 0, amount);
						    _videoRing.Write(_demuxMoveBuffer, 0, amount);
					    }
				    }
				    else if(_demuxTarget == DemuxTarget.Trash)
                        _demuxRing.Skip(amount);

                    if(!_noDemuxLimit)
				        _demuxRemaining -= amount;
			    }

                if (_demuxRemaining == 0)
                {
                    if (!_noDemuxLimit)
                        _demuxTarget = DemuxTarget.Parsing;
                }
		    }
	    }

        static int numFrames = 0;

	    void TryDecodeVideo()
	    {
		    if(_videoExhausted)
		    {
			    // Throw out any remaining video data
			    _videoRing.Skip(_videoRing.BytesStored);
			    return;
		    }

		    // Require enough preload that this is safe
		    if(_videoRing.BytesStored < VIDEO_BUFFER_SIZE && _parseState != PS.Exhausted)
			    return;

		    if(_videoDec == null)
		    {
                _videoDec = new VidStream(VIDEO_BUFFER_SIZE, this, _dispatch);
                _videoInitialized = false;
			    return;	// Don't do any more parsing this frame
		    }

            bool decResult = _videoDec.mpegVidRsrc(!_videoInitialized);
            _videoInitialized = true;
            if(_videoDec.FilmHasEnded)
			    _videoExhausted = true;
	    }

        void TryDecodeAudio()
	    {
		    // Require enough preload that this is safe
		    if(_audioRing.BytesStored < MP2DecodeState.FRAME_BUFFER_SIZE_BYTES + MP2DecodeState.FRAME_HEADER_SIZE_BYTES && _parseState != PS.Exhausted)
			    return;

            if (_audioRing.BytesStored < MP2DecodeState.FRAME_HEADER_SIZE_BYTES)
            {
                // Nothing left
                _audioExhausted = true;
                return;
            }

            _audioRing.Read(_audioData, 0, MP2DecodeState.FRAME_HEADER_SIZE_BYTES);
            if (_audioDec.ParseHeader(_audioData, 0, out _audioHeader))
            {
                _audioRing.Read(_audioData, 0, _audioDec.FrameBytes);
                _audioDec.DecodeFrame(_audioHeader, _audioData, 0, _audioSamples);
                _audioAvailable = true;
            }
	    }

        public void GetAudio(out short[] leftChannel, out short[] rightChannel, out uint sampleRate, out uint numChannels, out uint numSamples)
	    {
		    leftChannel = _audioSamples[0];
		    rightChannel = _audioSamples[1];
		    sampleRate = _audioDec.SampleRate;
		    numChannels = (_audioDec.InputStereo ? (uint)2 : (uint)1);
		    numSamples = MP2DecodeState.FRAME_NUM_SAMPLES;

		    _audioAvailable = false;

		    return;
	    }

        public bool AudioAvailable
        {
            get
            {
                return _audioAvailable;
            }
        }

	    public ImageInfo GetVideoFrame()
	    {
            _videoAvailable = false;
            return _videoImageInfo;
	    }

        public bool VideoAvailable
        {
            get
            {
                return this._videoAvailable;
            }
        }

        public void Work()
        {
            //using (PDB.Trace("Work"))
            {
                TryParsePackHeader();
                TryParseSystemHeader();
                TryParseStreams();
                TryDemux();
                TryDecodeAudio();
                TryDecodeVideo();
            }
        }

        public void Flush()
        {
            _flushing = true;
        }

        public bool IgnoreAudio { get; set; }
        public bool IgnoreVideo { get; set; }

        public bool Exhausted
        {
            get
            {
                return (_parseState == PS.Exhausted) && _audioExhausted && _videoExhausted;
            }
        }

        public void SetFrameSkip(int numFrames)
        {
            if (_videoDec != null)
                _videoDec.SetFrameSkip(numFrames);
        }

        public PerfDebugger VideoPDB
        {
            get
            {
                return _videoDec.PDB;
            }
        }

        RingBuffer IMpegVideoIO.Stream
        {
            get { return this._videoRing; }
        }

        void IMpegVideoIO.SetError(string p)
        {
        }

        void IMpegVideoIO.ExecuteDisplay(ImageInfo img)
        {
            _videoAvailable = true;
            _videoImageInfo = img;
        }
    }
}
